#Author: Mike Schladt - 2014
#Malware Intelligence Project MalwareSample Class
#This class defines structures require to track sample file information

import os
import re
import magic
import hashlib
import pydeep
import requests
from zipfile import ZipFile, BadZipfile

import logging
logging.basicConfig(filename='mip.log',level=logging.WARNING)

class MipSample : 
    
    
    def __init__(self,file = None):
        
        #if no file given, return empty container
        if file == None:
            return        

        self.errors = False
        self.timeout = 60
        self.file = file
        self.name = os.path.basename(self.file)
        self.size = os.path.getsize(self.file)
        self.ext = os.path.splitext(self.file)[-1].lower()
        self.filetype = magic.from_file(self.file)
        #add missing file extension
        if self.ext == "" : 
            self.ext = self.find_ext()  
            os.rename(self.file, self.file + self.ext)
            self.file = self.file + self.ext
        self.hashes()
        
    #method to calcule MD5,sha1,sha256,sha512,ssdeep    
    def hashes(self):
        #read contents of file
        data = open(self.file).read()
        
        #md5 section
        md5 = hashlib.md5()
        md5.update(data)
        self.md5 = md5.hexdigest()
        
        #sha1 
        sha1 = hashlib.sha1()
        sha1.update(data)
        self.sha1 = sha1.hexdigest()
        
        #sha256
        sha256 = hashlib.sha256()
        sha256.update(data)
        self.sha256 = sha256.hexdigest()
        
        #sha512
        sha512 = hashlib.sha512()
        sha512.update(data)
        self.sha512 = sha512.hexdigest()
        
        #ssdeep
        self.ssdeep = pydeep.hash_buf(data)
        
        del data #large data no longer needed
        return
    
    #method to return the best guess of file extension for passed file        
    def find_ext(self):
      
        #let's find some file extensions!!!
        if re.search('.*MS-DOS.*driver.*',self.filetype) :
            return ".dll"
            
        elif re.search('.*PE32.*',self.filetype) :
            return ".exe"
      
        elif re.search('.*PDF.*',self.filetype) :
            return ".pdf"
        
        elif re.search('.*HTML.*',self.filetype) :
            return ".html"

        elif re.search('.*PNG.*',self.filetype) :
            return ".png"

        elif re.search('.*GIF.*',self.filetype) :
            return ".gif"

        elif re.search('.*SVG.*',self.filetype) :
            return ".svg"

        elif re.search('.*JPEG.*',self.filetype) :
            return ".jpg"

        elif re.search('.*PHP.*',self.filetype) :
            return ".php"  

        elif re.search('.*Flash.*',self.filetype) :
            return ".swf"

        elif re.search('.*Word 2007\+.*',self.filetype) :
            return ".docx"

        elif re.search('.*Word,.*',self.filetype) :
            return ".doc"

        elif re.search('.*PowerPoint 2007\+.*',self.filetype) :
            return ".pptx"

        elif re.search('.*PowerPoint,.*',self.filetype) :
            return ".ppt"

        elif re.search('.*Excel 2007\+.*',self.filetype) :
            return ".xlsx"

        elif re.search('.*Excel,.*',self.filetype) :
            return ".xls" 

        elif re.search('.*shell script.*',self.filetype) :
            return ".sh" 

        elif re.search('.*batch file.*',self.filetype) :
            return ".bat"  
        
        #ZIPs and JARs
        elif re.search('.*Zip archive data.*',self.filetype) :
            with ZipFile(self.file, "r") as archive:
                zipinfos = archive.infolist()
                
                #search each file for .class or .java files
                isJar=False 
                for zipinfo in zipinfos:
                    #print zipinfo.filename
                    if re.search('(\.class)$',zipinfo.filename) \
                       or re.search('(\.java)$',zipinfo.filename):
                        isJar=True
                
                #apply either .jar or .zip extension
                if isJar:
                    return ".jar"
                else:
                    return ".zip" 
        
        #vbs js css - best guess
        elif re.search('.*ASCII text.*',self.filetype) :
            isVBS = False
            isJS = False
            isCSS = False
            #read contents of file
            data = open(self.file).read()
            #search for indicators
            if "Dim " in data : 
                isVBS = True
            elif "var " in data :
                isJS = True
            elif re.search('.*\{.*:.*\}.*',data) :
                isCSS = True
            
            del data # large data no longer needed    
            
            #assign extensions
            if isJS : 
                return ".js"
            elif isVBS :
              return ".vbs"
            elif isCSS :
                return ".css"
            else :
                return ".txt"
    
        else :
            return ""

    #method to connect and save sample attributes to db
    def store_db(self) : 
        import MySQLdb
        import _mysql_exceptions

        #username and password stored in external files
        mysql_usr = open(os.environ['HOME'] + "/.mysql_usr").read().strip()
        mysql_pwd = open(os.environ['HOME'] + "/.mysql_pwd").read().strip()
        
        #connect to database and create cursor
        db = MySQLdb.connect("localhost", mysql_usr, mysql_pwd, "mip")
        cursor = db.cursor()

        #create insert command
        cmd = ("INSERT INTO samples (name, size, filetype, md5, sha1, sha256, "
                "sha512, ssdeep) VALUES (%s,%s,%s,%s,%s,%s,%s,%s)")

        p1 = MySQLdb.escape_string(self.name)
        p2 = self.size
        p3 = MySQLdb.escape_string(self.filetype)
        p4 = MySQLdb.escape_string(self.md5)
        p5 = MySQLdb.escape_string(self.sha1)
        p6 = MySQLdb.escape_string(self.sha256)
        p7 = MySQLdb.escape_string(self.sha512)
        p8 = MySQLdb.escape_string(self.ssdeep)

        data = (p1,p2,p3,p4,p5,p6,p7,p8)
        
        try :        
            cursor.execute(cmd, data)
            self.sid = cursor.lastrowid

        #duplicate entry
        except _mysql_exceptions.IntegrityError as e :
            #parse key and value from error message
            value = str(e).split()[-4].split("'")[1]           
            key = str(e).split()[-1].split("'")[1]            
            #find sid of the orginal entry            
            cmd = "SELECT sid FROM samples WHERE " + key + "='" + value + "'"            
            
            try :             
                cursor.execute(cmd)
                self.sid = cursor.fetchone()[0]
            except Exception as e:
                print e
                logging.warning(e)

        except Exception as e:
            print "{0}, {1}".format(e, cmd)
            logging.warning("{0}, {1}".format(e, cmd))
    
        #close and clean up connection
        try:        
            cursor.close()
            db.close()
            del cursor
            del db
        except Exception as e:
            print e

    #module to submit sample to REST server 
    def submit(self) :

        request_url = "http://localhost:8090/tasks/create/file"
        sample_file = {'file': open(self.file, 'rb')}
        params = {'timeout': self.timeout, 'enforce_timeout': 1} #set params
        response = requests.post(request_url, files=sample_file, data=params);
        self.task_id = response.json()['task_id']
           
